import sqlite3
from datetime import datetime

class CashFlowDB:
    def __init__(self, db_path='investment_history.db'):
        """初始化数据库连接"""
        self.db_path = db_path
        self._create_table()
        
    def _create_table(self):
        """创建现金流水表"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                # 创建现金流水表
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS cash_flow (
                        flow_id INTEGER PRIMARY KEY AUTOINCREMENT,
                        user_id INTEGER NOT NULL,
                        asset_id INTEGER,
                        flow_type VARCHAR(16) NOT NULL,
                        amount DECIMAL(20,8) NOT NULL,
                        channel VARCHAR(64) NOT NULL,
                        fee DECIMAL(20,8) DEFAULT 0,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        note TEXT
                    )
                ''')
                conn.commit()
        except sqlite3.Error as e:
            print(f"创建cash_flow表时发生错误: {e}")
            raise
    
    def add_record(self, user_id, asset_id, flow_type, amount, channel, fee=0, note=None):
        """新增现金流水记录
        参数:
            user_id: 用户ID
            asset_id: 资产ID（法币时可为空）
            flow_type: 流水类型（deposit/withdraw/fiat_in/fiat_out）
            amount: 金额（正数为流入，负数为流出）
            channel: 渠道（银行转账、交易所、链上地址）
            fee: 手续费
            note: 备注
        返回:
            新增记录的ID
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    "INSERT INTO cash_flow (user_id, asset_id, flow_type, amount, channel, fee, note) VALUES (?, ?, ?, ?, ?, ?, ?)",
                    (user_id, asset_id, flow_type, amount, channel, fee, note)
                )
                conn.commit()
                return cursor.lastrowid
        except sqlite3.Error as e:
            print(f"新增现金流水记录时发生错误: {e}")
            raise
    
    def get_records_by_user(self, user_id, limit=100):
        """获取指定用户的所有现金流水记录
        参数:
            user_id: 用户ID
            limit: 限制返回记录数，默认100条
        返回:
            记录列表
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()
                cursor.execute(
                    "SELECT flow_id, asset_id, flow_type, amount, channel, fee, created_at, note FROM cash_flow WHERE user_id = ? ORDER BY created_at DESC LIMIT ?",
                    (user_id, limit)
                )
                rows = cursor.fetchall()
                return [dict(row) for row in rows]
        except sqlite3.Error as e:
            print(f"获取现金流水记录时发生错误: {e}")
            raise
    
    def get_record_by_id(self, flow_id, user_id):
        """获取指定用户的单条现金流水记录
        参数:
            flow_id: 流水ID
            user_id: 用户ID
        返回:
            记录字典，如果不存在或不属于该用户返回None
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()
                cursor.execute(
                    "SELECT flow_id, asset_id, flow_type, amount, channel, fee, created_at, note FROM cash_flow WHERE flow_id = ? AND user_id = ?",
                    (flow_id, user_id)
                )
                row = cursor.fetchone()
                return dict(row) if row else None
        except sqlite3.Error as e:
            print(f"获取现金流水记录时发生错误: {e}")
            raise